//
// Created by 邹迪凯 on 2021/7/17.
//

#ifndef LWPT_DIFFERENTIALPRIVACYASSIGNMENT_H
#define LWPT_DIFFERENTIALPRIVACYASSIGNMENT_H

#include <cstdlib>

#include "assignment.h"
#include "util.h"
#include "PublicRandomTaskAssignment.h"

class DPTAApplicationTupleChunk : public Chunk {
public:
    string taskId;
    string distanceWithNoise;

    DPTAApplicationTupleChunk(string taskId, string distanceWithNoise) : taskId(taskId),
                                                                         distanceWithNoise(distanceWithNoise) {}

    long size() const override {
        return taskId.size() + distanceWithNoise.size();
    }
};

class DPTAApplicationChunk : public Chunk {
public:
    string id;
    string privacyBudget;
    vector<DPTAApplicationTupleChunk> taskDistanceTuples;

    long size() const override {
        long ts = 0;
        for (auto &t:taskDistanceTuples) {
            ts += t.size();
        }
        return id.size() + privacyBudget.size() + ts + overhead();
    }

    long overhead() const {
        return 1 + taskDistanceTuples.size() - 1;
    }
};

class DPTAAssignment {
public:
    string id;
    double distance;
    double epsilon;

    DPTAAssignment(string id, double distance, double epsilon) : id(id), distance(distance), epsilon(epsilon) {}
};

class DPTAServer : public Server {
public:
    int maxApplicationCount;
    double taskRegionRedius;
    unordered_map<string, vector<DPTAAssignment> *> A;
    vector<Task> tasks;

    DPTAServer(int maxApplicationCount, double taskRegionRedius) : maxApplicationCount(maxApplicationCount),
                                                                   taskRegionRedius(taskRegionRedius) {
    }

    void prepareForAssignment() override {
        A.clear();
        for (const auto &ap : *applications) {
            auto a = static_cast<DPTAApplicationChunk *>(ap);
            for (const auto &apTup : (*a).taskDistanceTuples) {
                if (A.find(apTup.taskId) == A.end()) {
                    A[apTup.taskId] = new vector<DPTAAssignment>;
                }
                A[apTup.taskId]->emplace_back(a->id, stod(apTup.distanceWithNoise), stod(a->privacyBudget));
            }
        }
        for (auto &taskWorker:A) {
            auto ai = taskWorker.second;
            //bubble-sort
            for (int i = ai->size() - 1; i >= 0; i--) {
                for (int j = 0; j < i; j++) {
                    auto d = ai->at(j).distance - ai->at(j + 1).distance;
                    try {
                        if (integrateOverDoubleLaplace(d, ai->at(j).epsilon, ai->at(j + 1).epsilon) < 0.5) {
                            //di>dj
                            swap(*ai, j, j + 1);
                        }
                    } catch (...) {

                    }
                }
            }
        }
    }

    void assignTask() override {
        while (!A.empty()) {
            auto taskCandidateWorkers = *A.begin();
            if (taskCandidateWorkers.second->empty()) {
                A.erase(taskCandidateWorkers.first);
            }
            auto taskId = taskCandidateWorkers.first;
            auto firstCandidate = taskCandidateWorkers.second->front();
            vector<pair<string, vector<DPTAAssignment> *>> conflictedCandidates;
            //找到冲突的 fai 个 task-candidates
            for (auto &anotherTaskCandidateWorkers:A) {
                auto anotherFirstCandidate = anotherTaskCandidateWorkers.second->front();
                if (firstCandidate.id == anotherFirstCandidate.id) {
                    conflictedCandidates.emplace_back(anotherTaskCandidateWorkers);
                }
            }
            //bubble-sort
            for (int i = conflictedCandidates.size() - 1; i >= 0; i--) {
                for (int j = 0; j < i; j++) {
                    if (conflictedCandidates[j].second->size() == 1 ||
                        conflictedCandidates[j + 1].second->size() == 1) {
                        continue;
                    }
                    auto firstCandidate = conflictedCandidates[j].second->at(1);
                    auto secondCandidate = conflictedCandidates[j + 1].second->at(1);
                    auto d = firstCandidate.distance - secondCandidate.distance;
                    if (integrateOverDoubleLaplace(d, firstCandidate.epsilon, secondCandidate.epsilon) < 0.5) {
                        //di>dj
                        swap(conflictedCandidates, j, j + 1);
                    }
                }
            }
            for (int i = 0; i < conflictedCandidates.size(); i++) {
                if (conflictedCandidates[i].second->size() == 1) {
                    swap(conflictedCandidates, 0, i);
                    break;
                }
            }
            //分配任务
            auto smallestCandidates = conflictedCandidates.front();
            assignments.push_back(
                    new pair<string, string>(smallestCandidates.first, smallestCandidates.second->front().id));
            assignedTaskCount++;
            //任务分配出去了
            A.erase(smallestCandidates.first);
            for (auto c:conflictedCandidates) {
                swap(*c.second, 0, c.second->size() - 1);
                c.second->pop_back();
                if (c.second->empty()) {
                    A.erase(c.first);
                }
            }
        }

        for (const auto &item : A) {
            delete item.second;
        }
    }
};

bool tupleChunkComparator(const DPTAApplicationTupleChunk &t1, const DPTAApplicationTupleChunk &t2) {
    double d1 = stod(t1.distanceWithNoise), d2 = stod(t2.distanceWithNoise);
    return d1 < d2;
}

class DPTAWorker : public Worker {
public:
    double epsilon;

    DPTAWorker(string id, Location location, double epsilon) : Worker(id, location),
                                                               epsilon(epsilon) {}

    Chunk *apply(Server &server) override {
        //1. 从server的tasks中过滤出所有在同一个taskRegion中的task
        auto *dpaServer = static_cast<DPTAServer *>(&server);
        //2. 计算与task的distance，过滤出前n个
        vector<DPTAApplicationTupleChunk> candidateTuples;
        for (const auto &task : dpaServer->tasks) {
            double distance = earthDistance(stod(location.lat), stod(location.lon), stod(task.location.lat),
                                            stod(task.location.lon));
            if (distance < dpaServer->taskRegionRedius) {
                candidateTuples.emplace_back(task.id, to_string(distance));
            }
        }
        //排个序
        sort(candidateTuples.begin(), candidateTuples.end(), tupleChunkComparator);
        int limit = dpaServer->maxApplicationCount;
        auto *response = new DPTAApplicationChunk();
        response->id = id;
        response->privacyBudget = to_string(int(epsilon));
        for (auto &tup : candidateTuples) {
            //取前n个
            if (limit <= 0) {
                break;
            }
            //3. 添加噪声
            tup.distanceWithNoise = to_string(stod(tup.distanceWithNoise) + laplaceNoise(epsilon));
            response->taskDistanceTuples.push_back(tup);
            limit--;
        }
        //4. 发送chunk
        return response;
    }
};

class DPTATest : public Test {
public:
    int maxApplicationCount;
    double taskRegionRedius;
    double epsilon;

    DPTATest(DataSource *ds,
             int maxApplicationCount,
             double taskRegionRedius,
             double epsilon
    ) : Test(ds),
        maxApplicationCount(maxApplicationCount),
        taskRegionRedius(taskRegionRedius),
        epsilon(epsilon) {
    }

    Chunk *createTaskRegistration(Task &task) override {
        Chunk *c = new GenericRegisterChunk(task);
        static_cast<DPTAServer *>(server)->tasks.push_back(task);
        return c;
    }

    Worker *workerFactory(string id, Location location) override {
        if (epsilon == -1) {
            auto randomEpsilon = (double) (rand() / (double) RAND_MAX) * 5;
            return new DPTAWorker(id, location, randomEpsilon);
        }
        return new DPTAWorker(id, location, epsilon);
    }

    Server *serverFactory() override {
        return new DPTAServer(maxApplicationCount, taskRegionRedius);
    }

    string name() override {
        return string("DPTA(epsilon=").append(to_string(int(epsilon))).append(")");
    }
};

#endif //LWPT_DIFFERENTIALPRIVACYASSIGNMENT_H
